import pandas as pd
import re


def clean_data(input_path, output_path):
    # 读取原始数据
    df = pd.read_csv(input_path)

    # 清洗数据
    # 1. 处理"导演/主演"列，拆分为"导演"和"主演"
    # 使用正则表达式拆分“导演:”和“主演:”字段
    df[['导演', '主演']] = df['导演/主演'].str.extract(r'导演:\s*([^主演]+)\s*主演:\s*(.*)', expand=True)

    # 处理主演部分，拆分主演的多个演员
    df['主演'] = df['主演'].str.split(" / ")
    # 将主演列表转换为以逗号分隔的字符串
    df['主演'] = df['主演'].apply(lambda x: ', '.join(x) if isinstance(x, list) else x)

    # 2. 处理"上映年份/国别/类型"列，拆分为"年份"、"国别"和"类型"
    df[['年份', '国别', '类型']] = df['上映年份/国别/类型'].str.split("/", expand=True, n=2)

    # 3. 处理"评分"列，将评分转换为数字类型，缺失值替换为NaN
    df['评分'] = pd.to_numeric(df['评分'], errors='coerce')  # 评分转换为数字

    # 4. 处理"一句话影评"列，将其转换为字符串，缺失值替换为"暂无"
    df['一句话影评'] = df['一句话影评'].fillna("暂无")  # 填充缺失的影评

    # 5. 处理"年份"列，将年份中的非数字字符去除
    df['年份'] = df['年份'].apply(lambda x: re.sub(r'\D', '', str(x)))  # 去除非数字字符

    # 6. 删除不必要的列
    df = df.drop(columns=["导演/主演", "上映年份/国别/类型"])  # 删除原始的合并列

    # 7. 保存清洗后的数据
    df.to_csv(output_path, index=False, encoding="utf-8-sig")
    print(f"数据清洗完成，已保存到 {output_path}")


# 示例调用
input_path = './data/raw_data.csv'  # 原始爬取数据路径
output_path = './data/cleaned_data.csv'  # 清洗后数据路径

clean_data(input_path, output_path)
